import numpy as np
import cv2
import onnxruntime as ort
from onnx_utils import detect, detect_pose, nms
import os
import matplotlib.pyplot as plt
from yolo_decode import *


def sort_keypoints(keypoints):
    """
    Sort the keypoints by their position in the image.

    @param keypoints: 4行2列的关键点坐标，numpy数组，每行代表一个关键点的坐标(x,y)
    @return: 按左上、右上、右下、左下顺序排列的关键点坐标
    """
    index_left_up = np.argmin(np.sum(keypoints * np.array([1, 1]), axis=1))
    index_left_down = np.argmin(np.sum(keypoints * np.array([1, -1]), axis=1))
    index_right_down = np.argmax(np.sum(keypoints * np.array([1, 1]), axis=1))
    index_right_up = np.argmax(np.sum(keypoints * np.array([1, -1]), axis=1))
    keypoints_sorted = np.array(
        [
            keypoints[index_left_up],
            keypoints[index_right_up],
            keypoints[index_right_down],
            keypoints[index_left_down],
        ]
    )
    return keypoints_sorted


if __name__ == "__main__":
    model = ort.InferenceSession(
        r"C:\Users\Administrator\Desktop\ultralytics-action\训练结果\二维码\模型4\weights\best.onnx"
    )
    # img = cv2.imread(r"D:\desktopD\aidc_local\models\keypoint_detection\gitignored_dir\x1test\bb1.png")
    img_folder = r"D:\desktopD\aidc_cxftest\x1ScannerApp\app\src\main\assets\test_data\code2d_test_data"
    name_list = os.listdir(img_folder)
    for name in name_list:
        image = cv2.imread(os.path.join(img_folder, name))
        height, width, _ = image.shape
        length = max(height, width)
        imageB = np.zeros((256, 256, 3), np.float32)
        s = 256 / length

        resized_img = cv2.resize(image, (0, 0), fx=s, fy=s)
        h, w, c = resized_img.shape
        imageB[0:h, 0:w] = resized_img

        imageB = np.array([imageB.transpose(2, 0, 1)])
        imageB = np.mean(imageB, axis=1, keepdims=True)
        imageB /= 255

        output = model.run(None, {"input": imageB})
        out = output[0]

        res = nms(out)
        colors1 = [255, 0, 0]
        if res is not None:
            for rect in res:
                rect = rect[0].astype(int)
                cv2.rectangle(
                    resized_img,
                    (rect[0], rect[1]),
                    (rect[2], rect[3]),
                    colors1,
                    1,
                )

        plt.imshow(resized_img)
        plt.savefig(name)
